package jtg.parser;

import com.microsoft.z3.*;
import com.microsoft.z3.Context;
import com.microsoft.z3.Expr;
import soot.*;
import soot.jimple.*;
import soot.jimple.internal.*;

import java.lang.reflect.Field;
import java.util.*;

public class Z3Parser {

    public Context context;
    public Map<String, Expr> localExprs;
    public Map<String, Integer> assignList;
    public Solver solver = null;
    public FPRMExpr fpRM = null;

    public Z3Parser(Context ctx, Map<String, Expr> localExprs, Map<String, Integer> assignList) {
        this.context = ctx;
        this.localExprs = localExprs;
        this.assignList = assignList;
        this.fpRM = ctx.mkFPRoundNearestTiesToEven();
    }

    public Z3Parser(Z3Parser parser) {
        this.context = parser.context;
        this.localExprs = new HashMap<>(parser.localExprs);
        this.assignList = new HashMap<>(parser.assignList);
        this.solver = parser.solver;
        this.fpRM = parser.fpRM;
    }

    public Expr assignLocal(Local local, List<BoolExpr> cons) {
        Integer currentAssign = assignList.getOrDefault(local.toString(), -1);
        String tmpName = local + "_" + (currentAssign + 1);
        assignList.put(local.toString(), currentAssign + 1);
        Expr ret = fromLocal(local, cons);
        localExprs.put(tmpName, ret);
        return ret;
    }

    public  Expr assignInstanceField(Solver solver, InstanceFieldRef fieldRef, List<BoolExpr> cons) {
        Value ref = fieldRef.getBase();
        Expr ptr = solver.getModel().getConstInterp(fromSootExpr(ref, cons));
        if (ptr instanceof IntNum) {
            IntNum intNum = (IntNum) ptr;
            String className = fieldRef.getField().getDeclaringClass().getName();
            String fieldName = fieldRef.getField().getName();

            Local local = Jimple.v().newLocal(intNum + "_" + className + "." + fieldName, fieldRef.getField().getType());
            return assignLocal(local, cons);

        } else {
            System.err.println("Unsupported unconstrained pointer: " + ptr);
        }
        return null;
    }

    public BoolExpr getRangeConstraint(Context context, String name, long min, long max) {
        return context.mkAnd(context.mkGe(context.mkIntConst(name), context.mkInt(min)),
                context.mkLe(context.mkIntConst(name), context.mkInt(max)));
    }

    public Expr fromLocal(Local local, List<BoolExpr> constraints) {
        String name;
        if (assignList.get(local.getName()) != null)
            name = local.getName() + "_" + assignList.get(local.getName());
        else {
            System.err.println("fromLocal before assignLocal: " + local.getName());
            return null;
        }
        if (local.getType() instanceof soot.IntType) {
            // constraints.add(getRangeConstraint(context, name, Integer.MIN_VALUE, Integer.MAX_VALUE));
            return context.mkIntConst(name);
        }
        if (local.getType() instanceof soot.ShortType) {
            constraints.add(getRangeConstraint(context, name, Short.MIN_VALUE, Short.MAX_VALUE));
            return context.mkIntConst(name);
        }
        if (local.getType() instanceof soot.ByteType) {
            constraints.add(getRangeConstraint(context, name, Byte.MIN_VALUE, Byte.MAX_VALUE));
            return context.mkIntConst(name);
        }
        if (local.getType() instanceof soot.CharType) {
            constraints.add(getRangeConstraint(context, name, Character.MIN_VALUE, Character.MAX_VALUE));
            return context.mkIntConst(name);
        }
        if (local.getType() instanceof soot.LongType) {
            constraints.add(getRangeConstraint(context, name, Long.MIN_VALUE, Long.MAX_VALUE));
            return context.mkIntConst(name);
        }
        if (local.getType() instanceof soot.FloatType) {
            constraints.add(context.mkFPIsNormal(context.mkConst(name, context.mkFPSortSingle())));
            return context.mkConst(name, context.mkFPSortSingle());
        }
        if (local.getType() instanceof soot.DoubleType) {
            constraints.add(context.mkFPIsNormal(context.mkConst(name, context.mkFPSortDouble())));
            return context.mkConst(name, context.mkFPSortDouble());
        }
        if (local.getType() instanceof soot.BooleanType) {
            return context.mkBoolConst(name);
        }
        if (local.getType() instanceof soot.RefType) {
            Class<?> cls = null;
            try {
                cls = Class.forName(((soot.RefType) local.getType()).getClassName());
            } catch (ClassNotFoundException e) {
                System.err.println("Class not found: " + ((soot.RefType) local.getType()).getClassName());
                e.printStackTrace();
            }
            assert cls != null;
            // 若为String类型
            if (cls.equals(String.class)) {
                return context.mkConst(name, context.getStringSort());
            }
            if (cls.equals(StringBuilder.class)) {
                return context.mkConst(name, context.getStringSort());
            }
            if (cls.equals(java.io.PrintStream.class)) {
                return context.mkIntConst(name);
            }
            // 若为Enum子类
            if (cls.isEnum()) {
                // 获取所有枚举值
                Object[] enumConstants = cls.getEnumConstants();
                // 枚举值的个数
                int enumCount = enumConstants.length;
                constraints.add(getRangeConstraint(context, name, 0, enumCount - 1));
                return context.mkIntConst(name);
            }
            // 指针
            return context.mkIntConst(name);
            // 其他引用类型处理成对应大小的Seq<IntSort>
//            {
//                Field[] fields = cls.getDeclaredFields();
//                for (Field field : fields) {
//                    // 若非基本类型，跳过
//                    if (!field.getType().isPrimitive()) {
//                        System.err.println(cls.getName()+" Unsupported field type: " + field.getType());
//                        return null;
//                    }
//                }
//                int fieldCount = fields.length;
//                // constraints.add(context.mkEq(context.mkLength(context.mkConst(name, context.mkSeqSort(context.getIntSort()))),
//                //        context.mkInt(fieldCount)));
//                return context.mkConst(name, context.mkSeqSort(context.getIntSort()));
//            }
        }
        if (local.getType() instanceof soot.ArrayType) {

            ArrayType arrayType = (ArrayType) local.getType();

            // add length
            if (arrayType.numDimensions == 1) {
                String length_para = local.getName() + "_length";
                localExprs.put(length_para, context.mkIntConst(length_para));
            } else {
                System.out.println("Unsupported array dimension: " + arrayType.numDimensions);
                return null;
            }

            soot.Type elementType = arrayType.baseType;
            if (elementType instanceof soot.LongType || elementType instanceof soot.IntType || elementType instanceof soot.ShortType
                    || elementType instanceof soot.ByteType || elementType instanceof soot.CharType) {
                return context.mkArrayConst(name, context.getIntSort(), context.getIntSort());
            }
            if (elementType instanceof soot.FloatType || elementType instanceof soot.DoubleType) {
                return context.mkArrayConst(name, context.getIntSort(), context.getRealSort());
            }
            if (elementType instanceof soot.BooleanType) {
                return context.mkArrayConst(name, context.getIntSort(), context.getBoolSort());
            }
        }
        System.err.println("Unsupported local type: " + local.getType());
        return null;
    }

    public BoolExpr mkCastEq(Expr expr1, Expr expr2) {

        // System.out.println("leftType: " + leftType + ", rightType: " + rightType);
        if (expr1.isInt() && expr2.isBool()) {
            return context.mkEq(expr1, context.mkITE(expr2, context.mkInt(1), context.mkInt(0)));
        }
        if (expr1.isBool() && expr2.isInt()) {
            return context.mkEq(context.mkITE(expr1, context.mkInt(1), context.mkInt(0)), expr2);
        }
        return context.mkEq(expr1, expr2);
    }

    public boolean isFPType(soot.Type type) {
        return type instanceof soot.FloatType || type instanceof soot.DoubleType;
    }

    public Expr fromSootExpr(Value value) {
        return fromSootExpr(value, new ArrayList<>());
    }

    public Expr fromSootExpr(Value value, List<BoolExpr> constraints) {
        if (value instanceof Local) {
            Expr ret = localExprs.get(value + "_" + assignList.get(value.toString()));
            if (ret == null) {
                System.err.println("Local: " + value + "_" + assignList.get(value.toString()) + " is not found in localExprs");
            }
            return ret;
        }
        if (value instanceof IntConstant) {
            return context.mkInt(((IntConstant) value).value);
        }
        if (value instanceof LongConstant) {
            return context.mkInt(((LongConstant) value).value);
        }
        if (value instanceof FloatConstant) {
            return context.mkFP(((FloatConstant) value).value, context.mkFPSortSingle());
        }
        if (value instanceof DoubleConstant) {
            return context.mkFP(((DoubleConstant) value).value, context.mkFPSortDouble());
        }
        if (value instanceof StringConstant) {
            return context.mkString(((StringConstant) value).value);
        }
        if (value instanceof NullConstant) {
            return context.mkInt(0);
        }
        if (value instanceof JAddExpr) {
            if (isFPType(((JAddExpr) value).getOp1().getType())) {
                return context.mkFPAdd(fpRM, fromSootExpr(((JAddExpr) value).getOp1()),
                        fromSootExpr(((JAddExpr) value).getOp2()));
            }
            return context.mkAdd(fromSootExpr(((JAddExpr) value).getOp1()),
                    fromSootExpr(((JAddExpr) value).getOp2()));
        }
        if (value instanceof JSubExpr) {
            if (isFPType(((JSubExpr) value).getOp1().getType())) {
                return context.mkFPSub(fpRM, fromSootExpr(((JSubExpr) value).getOp1()),
                        fromSootExpr(((JSubExpr) value).getOp2()));
            }
            return context.mkSub(fromSootExpr(((JSubExpr) value).getOp1()),
                    fromSootExpr(((JSubExpr) value).getOp2()));
        }
        if (value instanceof JMulExpr) {
            if (isFPType(((JMulExpr) value).getOp1().getType())) {
                return context.mkFPMul(fpRM, fromSootExpr(((JMulExpr) value).getOp1()),
                        fromSootExpr(((JMulExpr) value).getOp2()));
            }
            return context.mkMul(fromSootExpr(((JMulExpr) value).getOp1()),
                    fromSootExpr(((JMulExpr) value).getOp2()));
        }
        if (value instanceof JDivExpr) {
            if (isFPType(((JDivExpr) value).getOp1().getType())) {
                return context.mkFPDiv(fpRM, fromSootExpr(((JDivExpr) value).getOp1()),
                        fromSootExpr(((JDivExpr) value).getOp2()));
            }
            return context.mkDiv(fromSootExpr(((JDivExpr) value).getOp1()),
                    fromSootExpr(((JDivExpr) value).getOp2()));
        }
        if (value instanceof JRemExpr) {
            if (isFPType(((JRemExpr) value).getOp1().getType())) {
                return context.mkFPRem(fromSootExpr(((JRemExpr) value).getOp1()),
                        fromSootExpr(((JRemExpr) value).getOp2()));
            }
            return context.mkRem(fromSootExpr(((JRemExpr) value).getOp1()),
                    fromSootExpr(((JRemExpr) value).getOp2()));
        }
        if (value instanceof JEqExpr) {
            if (isFPType(((JEqExpr) value).getOp1().getType())) {
                return context.mkFPEq(fromSootExpr(((JEqExpr) value).getOp1()),
                        fromSootExpr(((JEqExpr) value).getOp2()));
            }
            return mkCastEq(fromSootExpr(((JEqExpr) value).getOp1()),
                    fromSootExpr(((JEqExpr) value).getOp2()));
        }
        if (value instanceof JNeExpr) {
            return context.mkNot(context.mkEq(fromSootExpr(((JNeExpr) value).getOp1()),
                    fromSootExpr(((JNeExpr) value).getOp2())));
        }
        if (value instanceof JGeExpr) {
            return context.mkGe(fromSootExpr(((JGeExpr) value).getOp1()),
                    fromSootExpr(((JGeExpr) value).getOp2()));
        }
        if (value instanceof JGtExpr) {
            return context.mkGt(fromSootExpr(((JGtExpr) value).getOp1()),
                    fromSootExpr(((JGtExpr) value).getOp2()));
        }
        if (value instanceof JLeExpr) {
            return context.mkLe(fromSootExpr(((JLeExpr) value).getOp1()),
                    fromSootExpr(((JLeExpr) value).getOp2()));
        }
        if (value instanceof JLtExpr) {
            return context.mkLt(fromSootExpr(((JLtExpr) value).getOp1()),
                    fromSootExpr(((JLtExpr) value).getOp2()));
        }
        if (value instanceof JAndExpr) {
            return context.mkAnd(fromSootExpr(((JAndExpr) value).getOp1()),
                    fromSootExpr(((JAndExpr) value).getOp2()));
        }
        if (value instanceof JOrExpr) {
            return context.mkOr(fromSootExpr(((JOrExpr) value).getOp1()),
                    fromSootExpr(((JOrExpr) value).getOp2()));
        }
        if (value instanceof JXorExpr) {
            return context.mkXor(fromSootExpr(((JXorExpr) value).getOp1()),
                    fromSootExpr(((JXorExpr) value).getOp2()));
        }
        if (value instanceof JShlExpr) {
            return context.mkBVSHL(fromSootExpr(((JShlExpr) value).getOp1()),
                    fromSootExpr(((JShlExpr) value).getOp2()));
        }
        if (value instanceof JShrExpr) {
            return context.mkBVASHR(fromSootExpr(((JShrExpr) value).getOp1()),
                    fromSootExpr(((JShrExpr) value).getOp2()));
        }
        if (value instanceof JUshrExpr) {
            return context.mkBVLSHR(fromSootExpr(((JUshrExpr) value).getOp1()),
                    fromSootExpr(((JUshrExpr) value).getOp2()));
        }
        if (value instanceof JNegExpr) {
            return context.mkUnaryMinus(fromSootExpr(((JNegExpr) value).getOp()));
        }
        // 浮点数
        if (value instanceof JCastExpr) {
             System.out.println("JCastExpr: " + ((JCastExpr) value).getCastType());
        }
        if (value instanceof JCmplExpr) {
            FPExpr expr1 = (FPExpr) fromSootExpr(((JCmplExpr) value).getOp1());
            FPExpr expr2 = (FPExpr) fromSootExpr(((JCmplExpr) value).getOp2());
            return context.mkITE(context.mkFPEq(expr1, expr2), context.mkInt(0),
                    context.mkITE(context.mkFPLt(expr1, expr2), context.mkInt(-1),
                            context.mkInt(1)));
        }
        if (value instanceof JCmpgExpr) {
            FPExpr expr1 = (FPExpr) fromSootExpr(((JCmpgExpr) value).getOp1());
            FPExpr expr2 = (FPExpr) fromSootExpr(((JCmpgExpr) value).getOp2());
            return context.mkITE(context.mkFPEq(expr1, expr2), context.mkInt(0),
                    context.mkITE(context.mkFPLt(expr1, expr2), context.mkInt(-1),
                            context.mkInt(1)));
        }
        if (value instanceof JNewArrayExpr) {
            // baseType
            Type baseType = ((JNewArrayExpr) value).getBaseType();
            if (baseType instanceof soot.IntType || baseType instanceof soot.ShortType
                    || baseType instanceof soot.ByteType || baseType instanceof soot.CharType
                    || baseType instanceof soot.LongType) {
                return context.mkConstArray(context.getIntSort(), context.mkInt(0));
            }
            if (baseType instanceof soot.FloatType) {
                return context.mkConstArray(context.getIntSort(), context.mkFP(0.0f, context.mkFPSortSingle()));
            }
            if (baseType instanceof soot.DoubleType) {
                return context.mkConstArray(context.getIntSort(), context.mkFP(0.0, context.mkFPSortDouble()));
            }
            if (baseType instanceof soot.BooleanType) {
                return context.mkConstArray(context.getIntSort(), context.mkFalse());
            }
            System.err.println("Unsupported baseType: " + baseType);
        }
        if (value instanceof JArrayRef) {
            return context.mkSelect(fromSootExpr(((JArrayRef) value).getBase()),
                    fromSootExpr(((JArrayRef) value).getIndex()));
        }
        if (value instanceof JLengthExpr) {
//            System.out.println("JLengthExpr: " + ((JLengthExpr) value).getOp().toString() + "_length");
            return localExprs.get(((JLengthExpr) value).getOp().toString() + "_length");
        }
        if (value instanceof StaticFieldRef) {
            SootField field = ((StaticFieldRef) value).getField();
            String className = field.getDeclaringClass().getName();
            String filedName = field.getName();
            Class<?> cls = null;
            try {
                cls = Class.forName(className);
            } catch (ClassNotFoundException e) {
                System.err.println("Class not found: " + className);
                e.printStackTrace();
            }
            assert cls != null;
            // System.out.println("StaticFieldRef: " + className + "." + filedName);
            if (cls.isEnum()) {
                Object[] enumConstants = cls.getEnumConstants();
                int enumCount = enumConstants.length;
                for (int i = 0; i < enumCount; i++) {
                    if (enumConstants[i].toString().equals(filedName)) {
                        return context.mkInt(i);
                    }
                }
            }
            if (cls.equals(java.lang.System.class)) {
                return context.mkInt(0);
            }
        }
        if (value instanceof VirtualInvokeExpr) {
            SootMethod method = ((VirtualInvokeExpr) value).getMethod();
            String className = method.getDeclaringClass().getName();
            String methodName = method.getName();
            Value base = ((VirtualInvokeExpr) value).getBase();
//            System.out.println("VirtualInvokeExpr: " + className + "." + methodName + "(" + base + ")");
            if (className.equals("java.lang.String")) {
                if (methodName.equals("equals")) {
                    return context.mkEq(fromSootExpr(base), fromSootExpr(((VirtualInvokeExpr) value).getArg(0)));
                }
                if (methodName.equals("length")) {
                    return context.mkLength(fromSootExpr(base));
                }
                if (methodName.equals("charAt")) {
                    return context.mkSelect(fromSootExpr(base), fromSootExpr(((VirtualInvokeExpr) value).getArg(0)));
                }
            }
            if (className.equals("java.lang.StringBuilder")) {
                if (methodName.equals("append")) {
//                    System.out.println("append: " + base + ", " + ((VirtualInvokeExpr) value).getArg(0));
                    return context.mkConcat(
                            (Expr<SeqSort<BitVecSort>>)fromSootExpr(base),
                            (Expr<SeqSort<BitVecSort>>)fromSootExpr(((VirtualInvokeExpr) value).getArg(0)));
                }
                if (methodName.equals("toString")) {
                    return fromSootExpr(base);
                }
            }
        }
        if (value instanceof JStaticInvokeExpr) {
            SootMethod method = ((JStaticInvokeExpr) value).getMethod();
            String className = method.getDeclaringClass().getName();
            String methodName = method.getName();
            if (methodName.equals("toString")) {
                Expr expr = fromSootExpr(((JStaticInvokeExpr) value).getArg(0));
                if (expr.isInt()) {
                    return context.intToString(expr);
                }
            }
        }
        if (value instanceof JNewExpr) {
            String className = ((JNewExpr) value).getBaseType().toString();
            if (className.equals("java.lang.StringBuilder")) {
                return context.mkString("");
            }
            // 随机数指针
            return context.mkInt(new Random().nextInt());
        }
        if (value instanceof InstanceFieldRef) {
            InstanceFieldRef fieldRef = (InstanceFieldRef) value;
            Value ref = fieldRef.getBase();
            Expr ptr = fromSootExpr(ref, constraints);
            if (solver != null) {
                ptr = solver.getModel().getConstInterp(ptr);
            }
            if (ptr instanceof IntNum) {
                IntNum intNum = (IntNum) ptr;
                String className = fieldRef.getField().getDeclaringClass().getName();
                String fieldName = fieldRef.getField().getName();
                Local local = Jimple.v().newLocal(intNum + "_" + className + "." + fieldName, fieldRef.getField().getType());
                return fromSootExpr(local, constraints);
            } else {
                System.err.println("Unsupported unconstrained pointer: " + ptr);
            }
        }
        System.out.println("Unsupported value type: " + value.getClass().getName());
        return null;
    }
}
